# aem.two_d.transient.head_well --
#
# Copyright (c) 2000, V.Kelson and M.Bakker (steady well)
#           (c) 2001, Willem Jan Zaadnoordijk (transient well)
#           (c) 2004, Willem Jan Zaadnoordijk (tr. head well)

from types import *

import copy

import geometry
XY = geometry.XY
SearchableList = geometry.two_d.collections.SearchableList

import base

import aem.container
RowList = aem.container.RowList
ColumnList = aem.container.ColumnList

import aem.matrix



class HeadWell(geometry.two_d.Circle,base.Element):
	'''Class for transient wells

	Python call: HeadWell(xc,yc,radius,tHlist=[geometry.XY(0,0)],
	                  parent=None,controlPoints=None)

	Attributes:

	  - xc,yc: location of well (inherited from Circle)
	  - radius: radius of well (inherited from Circle)
	  - tHlist:  list of pairs (time, head) with time in ascending order
	  - parent: parent of the element (inherited from Element)
	  - parameters: ColumnList of strength parameters, set to empty list (inherited from Element)
	  - solver: solver (inherited from Element)
	  - controlPoints: list of control points, optional, set to right side of well radius
	  - elementList: list of transient elements ("UnitStepWell"s)
	  - initialSteadyState: list of initial steady element(s)
	    this element is not included in 'parameters'and 'elementList'
	    and does not appear in the influence functions

	Overloaded functions from Element:
		__init__, getMatrixCoefficients, getMatrixContribution,
		check, potentialInfluence, dischargeInfluence
	'''

	def __init__(self,xc,yc,radius,tHlist=[geometry.XY(0,0)],parent=None,controlPoints=None):
		'''Constructor'''
		geometry.two_d.Circle.__init__(self,xc,yc,radius)
		base.Element.__init__(self,parent)
		self.elementList = SearchableList()
		self.__numberOfParameters = 0
		self.parameters=ColumnList([])
		self.initialH  = None
		self.tlastsolve= -1
		self.specifiedH= []
		self.specifiedt= []
		self.givenElement = 0
		if controlPoints:
			self.controlPoints = copy.copy(controlPoints)
		else:
			self.controlPoints = [ self.center + geometry.XY(self.radius,0) ]
		for partim in tHlist :
			if partim.x <= 0 :
				if self.tlastsolve < 0 :
					# no initial steady state has been specified yet
					self.initialH = partim.y
					self.initialSteadyState = aem.two_d.single.HeadWell( \
						xc,yc,radius,specifiedHead=self.initialH,         \
						parent=parent.initialSteadyState,                 \
						controlPoints=self.controlPoints)
					# initial steady state is starting head for transient calculation
					self.tlastsolve = 0
					self.specifiedH.append(partim.y)
					self.specifiedt.append(partim.x)
			else :
				self.specifiedH.append(partim.y)
				self.specifiedt.append(partim.x)
				if self.tlastsolve < 0 : selftlastsolve = partim.x

	def __repr__(self):
#		# Parent not yet included
#		string = 'Well' + str((self.center.x,self.center.y,self.radius))
#		print string
#		for el in self.elementList:
#			string = string + el.__repr__()
#			print string
#		return string
		return 'Trans.Well{'    \
		   + str((self.center.x,self.center.y,self.radius, \
		         self.initialH,self.specifiedt,self.specifiedH)) \
		   + '}'

	def addElement(self,el):
		self.elementList.append(el)

	def getMatrixCoefficients(self,fa,ctlPt,t):
		# return value of the FA influence function at point CTLpT and time T
		# for the last UnitStepWell of this element

		if t > self.tlastsolve :
			for i in range(len(self.specifiedt)) :
				if t > self.specifiedt[i] :
					# determine start time connected to this solve time and
					# add UnitStepWell if necessary
					# - same code as in getMatrixContribution, because
					#   it is unknown which function gets called first
					nelem=len(self.elementList)
					if nelem==0:
						tlast=self.tlastsolve-1
					else:
						elemlast = self.elementList[len(self.elementList)-1]
						tlast=elemlast.t
					if tlast<self.tlastsolve:
						aem.two_d.transient.UnitStepWell(         \
							self.tlastsolve,self)
						self.parameters.append(0)
						self.__numberOfParameters = self.__numberOfParameters+1
					# start time has been determined (and UnitStepWell added)
					# - same code as in getMatrixCoefficients, because
					#   it is unknown which function gets called first
		if len(self.elementList)>0:
			if self.isInArea(ctlPt):
				xyw = XY( self.radius , 0 )
			else:
				xyw = XY( ctlPt.x-self.zc.real , ctlPt.y-self.zc.imag )
			return fa(self.elementList[(len(self.elementList)-1)],xyw,t)
		else:
			return 0

	def getMatrixContribution(self,elementList,t):
		# Returns a tuple (mc,mr) where mc is a list of MatrixColumn objects
		# and mr is a list of MatrixRow objects
		# specifying the condition at time t of a new
		# UnitStepLineSink starting at the last time of solving
		rows = []
		columns = []
		specifiedHead= None
		active = 0
		if t > self.tlastsolve :
			for i in range(len(self.specifiedt)) :
				if t > self.specifiedt[i] :
					specifiedHead = self.specifiedH[i]
					active = 1
		if active==1 :
			# determine start time connected to this solve time and
			# add UnitStepWell if necessary
			# - same code as in getMatrixContribution, because
			#   it is unknown which function gets called first
			nelem=len(self.elementList)
			if nelem==0:
				tlast=self.tlastsolve-1
			else:
				elemlast = self.elementList[len(self.elementList)-1]
				tlast=elemlast.t
			if tlast<self.tlastsolve:
				tstart = max(self.tlastsolve,self.specifiedt[0])
				aem.two_d.transient.UnitStepWell(       \
					tstart,self)
				self.parameters.append(0)
				self.__numberOfParameters = self.__numberOfParameters+1
			# start time has been determined (and UnitStepWell added)
			# - same code as in getMatrixContribution, because
			#   it is unknown which function gets called first
			columns.append(aem.matrix.MatrixColumn(self,len(self.elementList)-1))
			for cp in self.controlPoints:
				self.alpha=self.parent.aquiferDiffusivity(cp,t)
				mr = aem.matrix.MatrixRow()
				for e in elementList:
					mr.A = mr.A + e.getMatrixCoefficients(lambda z,p,t:z.potentialInfluence(p,t),cp,t)
					mr.b = self.headToPotential(cp,specifiedHead) - \
						            self.potential(cp,t)
					addA=e.getMatrixCoefficients(lambda z,p,t:z.potentialInfluence(p,t),cp,t)
			rows.append(mr)
		self.tlastsolve = t
		return ( rows,columns )

	def check(self,t):
		# Checks the boundary condition(s)
		rv = []
		active = 0
		for i in range(len(self.specifiedt)) :
			if t > self.specifiedt[i] :
				specifiedHead= self.specifiedH[i]
				active = 1
		if active==1 :
			for cp in self.controlPoints:
				rv.append(self.head(cp,t) - specifiedHead)
		return rv

	def potentialInfluence(self,xy,t):
		assert isinstance(xy,XY),"XY point is required"
		pot=[]
		self.alpha=self.parent.aquiferDiffusivity(xy,t)
		if self.isInArea(xy):
			xyw = XY( self.radius , 0 )
		else:
			xyw = XY( xy.x-self.zc.real , xy.y-self.zc.imag )
		for el in self.elementList:
			pot = pot + el.potentialInfluence(xyw,t)
		return RowList( pot )

	def dischargeInfluence(self,xy,t):
		assert isinstance(xy,XY),"XY point is required"
		dis=[]
		self.alpha=self.parent.aquiferDiffusivity(xy,t)
		if self.isInArea(xy):
			xyw = XY( self.radius,0)
		else:
			xyw = XY( xy.x-self.zc.real , xy.y-self.zc.imag )
		for el in self.elementList:
			dis = dis+el.dischargeInfluence(xyw,t)
		return RowList( dis )
